#!/usr/bin/python3
'''
这个文件只包含一个通信类 Communication
负责实现与服务器的通信，包含一些加解密操作与cloudcrypto相配合
'''
import socket
import struct
import json
import os
from cloudcrypto import *

INIT = 0

class Communication:
    '''
    用于实现和服务器通信，主要操作如下：
    1.register 2.login 3.upload 4.download 5.delete 6.updatedir 7.share 8.getshare
    '''
    address_family = socket.AF_INET
    socket_type = socket.SOCK_STREAM
    max_packet_size = 8192
    coding = 'utf-8'
    request_queue_size = 10

    def __init__(self, server_address, connect=True):
        
        self.server_address = server_address
        self.client_socket = socket.socket(self.address_family, self.socket_type)
        if connect:
            try:
                self.client_socket.connect(self.server_address)
            except:
                self.client_socket.close()
                raise        

    @staticmethod
    def util_packhead(head_content):
        '''
        工具函数，pack协议头部信息，为了保证协议的规范性，保证head是字典
        返回封装后的字节流，前面是内容长度后面是内容
        '''
        head_json = json.dumps(head_content)
        head_bytes = bytes(head_json, encoding=Communication.coding)
        return struct.pack('i', len(head_bytes)) + head_bytes

    @staticmethod
    def util_unpackhead(conn):
        '''
        工具函数，获取协议头部信息，为了保证协议的规范性，保证头部都是json化的字典
        '''
        head_len = struct.unpack('i', conn.recv(4))[0]
        head_json = conn.recv(head_len).decode(Communication.coding)
        head_content = json.loads(head_json)
        return head_content

    @staticmethod
    def util_packcode(code):
        '''
        工具函数，code -> pcode
        返回pack后的字节流
        '''
        return struct.pack('i', code)

    @staticmethod
    def util_unpackcode(pcode):
        '''
        工具函数，pcode -> code
        返回unpack后的code
        '''
        return struct.unpack('i', pcode)[0]

    @staticmethod
    def util_recvdata(conn, filesize):
        '''
        接收文件数据，输入是连接和文件的大小，因为返回的是bytes类型
        这里一般用作接收文件，根据filesize，得到的内容全部返回
        '''
        recv_size = 0
        recv_datas = bytes()
        while recv_size < filesize:
            recv_data = conn.recv(Communication.max_packet_size)
            recv_size += len(recv_data) 
            recv_datas = recv_datas + recv_data # 注意顺序
            # print('recvsize:%s filesize:%s' % (recv_size, filesize)) 这一条如果后续要弄进度条会非常有用，暂时没用，因为我不打算弄进度条了，哈哈^_^
        return recv_datas

    def register(self):
        '''
        注册只能先搁置了，现在有的用户直接保存在服务器代码中，是alice，bob和lyp
        '''
        print('This is register!')

    def login(self, *args):
        '''
        发送登录信息，如果成功登录则返回1否则返回0
        登录成功后，还需要从服务器获得目录文件directory.txt，放到./用户名/etc/中，供后续初始化读取
        ''' 
        user_info = args[0]

        # 构造并发送pcode和包头
        pcode = Communication.util_packcode(2)
        head = Communication.util_packhead(user_info)
        self.client_socket.send(pcode + head)
        
        # 接收返回码和下载directory.txt
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            # 接收directory.txt
            head_content = Communication.util_unpackhead(self.client_socket)
            if 'filesize' in head_content:
                filesize = head_content['filesize']
                recv_datas = Communication.util_recvdata(self.client_socket, filesize)
            else:
                return 0 # 接收错误的包头信息，但后面不清理recvbuffer必然会出错，但可能性比较小，就不管了 20211002

            # 初始化路径
            self.basepath = './' + user_info['account'] # 基路径
            self.client_data_path = self.basepath + '/data'
            self.client_etc_path = self.basepath + '/etc'

            # 将数据写入到etc/directory.txt中
            with open(self.client_etc_path+'/directory.txt', 'wb') as f:
                f.write(recv_datas)

            # 注意第一次时需要生成密钥对，使用cloudcrypto的cloud_pre_genkey函数，本来这里写到注册会比较好
            if INIT == 1:
                cloud_pre_genkey(self.client_etc_path)

            # 读取并还原签名密钥对
            self.my_signing_key, self.my_verifying_key = cloud_pre_restore_signingkeypair(self.client_etc_path)
            self.my_signer = Signer(self.my_signing_key)

            # 读取并还原授权密钥对
            self.my_secret_key, self.my_public_key = cloud_pre_restore_delegatingkeypair(self.client_etc_path)
            return 1
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            print(recv_buffer.decode(Communication.coding))
            return 0

    def upload(self, *args):
        '''
        如果上传成功，则返回1，否则返回0
        '''
        file_src = args[0]
        file_dst = args[1]
        if not os.path.isfile(file_src):  # 判断路径是否为文件
            print('file:{} is not exists'.format(file_src))
            return 0

        # 构造并发送包头的内容
        filesize = os.path.getsize(file_src) 
        filename = cloud_sha256(file_dst) # file_dst like '/3.txt'
        file_info = {'filename': filename, 'filesize': filesize}
        pcode = Communication.util_packcode(3)
        head = Communication.util_packhead(file_info)
        self.client_socket.send(pcode + head)
        
        # 发送部分
        # 读取明文
        with open(file_src, 'rb') as f:
            plaintext = f.read() 
        # 加密，生成capsule和ciphertext
        capsule, ciphertext = encrypt(self.my_public_key, plaintext)
        # 发送文件
        send_size = 0 
        send_datas = capsule.__bytes__() + ciphertext # data是不可数名词，但我这里要分开发包，不过这里发文件没有之前那么方便，不过也还行，自己构造索引
        while (send_size < len(send_datas)):
            send_data = send_datas[send_size:send_size+Communication.max_packet_size]
            self.client_socket.send(send_data)
            send_size += len(send_data)

        # 接收返回信息
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            return 1
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            print(recv_buffer.decode(Communication.coding))
            return 0

    def download(self, *args):
        '''
        如果下载成功，则返回1，否则返回0
        file_src(cloud) --down--> file_dst(local)
        '''
        file_src = args[0] # /1/4 是一个文件
        file_dst = args[1] # './data' 只是一个目录

        # 发送请求包，告知服务器需要下载的文件
        filename = cloud_sha256(file_src)
        file_info = {'filename': filename}
        pcode = Communication.util_packcode(4)
        head = Communication.util_packhead(file_info)
        self.client_socket.send(pcode + head)

        # 接收返回码和下载文件
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            head_content = Communication.util_unpackhead(self.client_socket)
            if 'filesize' in head_content:
                filesize = head_content['filesize']

                # 接收文件数据
                recv_datas = Communication.util_recvdata(self.client_socket, filesize)
                
                # 还原capsule，ciphertext并解密
                capsule = Capsule._from_exact_bytes(recv_datas[0:Capsule.serialized_size()])
                ciphertext = recv_datas[Capsule.serialized_size():]
                cleartext = decrypt_original(self.my_secret_key, capsule, ciphertext)
                
                # 写入文件
                filepath = os.path.normpath(os.path.join(file_dst, os.path.basename(file_src)))
                with open(filepath, 'wb') as f:
                    f.write(cleartext)                  
                return 1
            else:
                return 0 # 接收错误的包头信息，但后面不清理recvbuffer必然会出错，但可能性比较小，就不管了 20211002
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            print(recv_buffer.decode(Communication.coding))
            return 0

    def delete(self, *args):
        '''
        如果删除成功，则返回1，否则返回0
        delete file_dst
        '''
        file_dst = args[0] # /1/2.txt类似这样的

        # 发送请求包，告知服务器需要删除的文件
        filename = cloud_sha256(file_dst) # file_dst like '/3.txt'
        file_info = {'filename': filename}
        pcode = Communication.util_packcode(5)
        head = Communication.util_packhead(file_info)
        self.client_socket.send(pcode + head)

        # 接收返回信息
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            return 1
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            print(recv_buffer.decode(Communication.coding))
            return 0

    def updatedir(self, *args):
        '''
        更新目录，和更新公钥（为了方便我决定把公钥也通过这个发送给服务器）
        '''
        directory = args[0]
        directory_bytes = directory.encode(Communication.coding)

        # 公钥一并发给服务器
        with open(self.client_etc_path+'/delegating_key.txt', 'rb') as f:
            f.seek(32) # 移动文件读取指针到指定位置，移动到公钥的位置
            public_key_bytes = f.read(33)
        dirpublic_bytes = public_key_bytes + directory_bytes # public_key不是bytes形式，是一个类

        # 发送包头
        file_info = {'filesize': len(dirpublic_bytes)}
        pcode = Communication.util_packcode(6)
        head = Communication.util_packhead(file_info)
        self.client_socket.send(pcode + head)
        
        # 发送公钥和目录bytes
        self.client_socket.send(dirpublic_bytes)
        
        # 接收返回信息
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            return 1
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            print(recv_buffer.decode(Communication.coding))
            return 0

    def share(self, *args):
        '''
        分享文件，如果分享成功返回1，否则返回0
        '''
        account = args[0]
        file_src = args[1]

        # 发送请求包，告知服务器待分享用户的账号和需要分享的文件，
        filename = cloud_sha256(file_src) # file_dst like '/3.txt'
        share_info = {'account': account, 'filename': filename}
        pcode = Communication.util_packcode(7)
        head = Communication.util_packhead(share_info)
        self.client_socket.send(pcode + head)

        # 接收返回信息，得到被分享者的公钥
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len) 
            bob_public_key_bytes = recv_buffer
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            return 0, recv_buffer.decode(Communication.coding) # 错误返回错误信息

        # 计算kfrags，并发送给服务器。因为我是减法使用umbral，因此只要threshold和shares为1即可
        bob_public_key = PublicKey._from_exact_bytes(bob_public_key_bytes)
        kfrags = generate_kfrags(delegating_sk=self.my_secret_key,
                                 receiving_pk=bob_public_key,
                                 signer=self.my_signer,
                                 threshold=1,
                                 shares=1)
        kfrag_bytes = kfrags[0].__bytes__() # 只有一个
        
        # # 保留kfrag
        # with open(self.client_etc_path+'/kfrag_bytes.txt', 'wb') as f:
        #     f.write(kfrag_bytes)
        kfrag_info = {'kfragsize': VerifiedKeyFrag.serialized_size()}
        head = Communication.util_packhead(kfrag_info)
        self.client_socket.send(head)
        self.client_socket.send(kfrag_bytes)

        # 接收返回信息，得到分享链接
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            return 1, filename
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            return 0, recv_buffer.decode(Communication.coding) # 错误返回错误信息

    def getshare(self, *args):
        '''
        获取分享的文件
        '''
        saccount = args[0]
        filename = args[1]

        # 发送请求包
        share_info = {'account': saccount, 'filename': filename}
        pcode = Communication.util_packcode(8)
        head = Communication.util_packhead(share_info)
        self.client_socket.send(pcode + head)

        # 接收返回码和下载文件
        recv_buffer = self.client_socket.recv(4)
        status = Communication.util_unpackcode(recv_buffer) # status表示返回码
        if status == 1:
            # 接收分享者公钥并还原
            publickey_len = PublicKey.serialized_size()
            publickey_bytes = self.client_socket.recv(publickey_len)
            publickey = PublicKey._from_exact_bytes(publickey_bytes)

            # 接收kfrag_bytes并还原
            kfrag_len = VerifiedKeyFrag.serialized_size()
            kfrag_bytes = self.client_socket.recv(kfrag_len)
            kfrag = VerifiedKeyFrag.from_verified_bytes(kfrag_bytes)

            head_content = Communication.util_unpackhead(self.client_socket)
            if 'filesize' in head_content:
                filesize = head_content['filesize']
   
                # 接收文件数据
                recv_datas = Communication.util_recvdata(self.client_socket, filesize)
                
                # 获取并还原capsule，ciphertext  
                capsule_len = Capsule.serialized_size()
                capsule_bytes = recv_datas[0:capsule_len]
                capsule = Capsule._from_exact_bytes(capsule_bytes)
                ciphertext = recv_datas[capsule_len:]

                # 根据kfrag重加密生成cfrag
                cfrag = reencrypt(capsule=capsule, kfrag=kfrag)
                cfrags = [cfrag]
                
                cleartext = decrypt_reencrypted(receiving_sk=self.my_secret_key,
                                                delegating_pk=publickey,
                                                capsule=capsule,
                                                verified_cfrags=cfrags,
                                                ciphertext=ciphertext)

                # 写入文件
                filepath = os.path.normpath(os.path.join(self.client_data_path, filename)) 
                with open(filepath, 'wb') as f:
                    f.write(cleartext)         
                return 1, filename
            else:
                return 0,'接收错误的包头信息' # 接收错误的包头信息，但后面不清理recvbuffer必然会出错，但可能性比较小，就不管了 20211002
        else:
            recv_buffer = self.client_socket.recv(4)
            head_len = struct.unpack('i', recv_buffer)[0]
            recv_buffer = self.client_socket.recv(head_len)
            return 0, recv_buffer.decode(Communication.coding)

    def close_connection(self):
        self.client_socket.close()

# client = Communication(('192.168.59.133', 8080))
# client.manage_operates()